# ==============================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         posenet_train.py
# Date:           2025-05-24
# Description:    本模块实现了一个基于深度学习的姿态估计模型 PoseNet，用于从灰度图像中回归旋转向量 (rvec)。
#                 模型结构包含一个 VGG 风格的卷积主干网络，并结合残差模块提升特征表达能力。
#                 使用 Rodrigues 公式将旋转向量转换为旋转矩阵，并定义 Geodesic Loss 作为训练损失函数，
#                 以更准确地衡量预测姿态与真实姿态之间的角度误差。
#                 此外，模型训练过程中采用了 AdamW 优化器、早停机制（Early Stopping）和学习率自适应调整策略（ReduceLROnPlateau），
#                 并计算了 MAE、RMSE 和 Pearson 相关系数等多种评估指标来监控训练效果。
# ==============================================================================

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
import numpy as np
import time
import os
import cv2
import csv
from tqdm import tqdm
from torch.optim.lr_scheduler import OneCycleLR, CosineAnnealingLR, LinearLR, SequentialLR, ReduceLROnPlateau


# ==========================================================
#           辅助函数：Rodrigues 转换和 Geodesic Loss
# ==========================================================
def batch_rodrigues(rvecs):
    angle = torch.norm(rvecs, dim=1, keepdim=True)
    small_angle = angle < 1e-6
    axis = rvecs / (angle + 1e-8)
    axis = axis * (~small_angle).float()
    angle = angle * (~small_angle).float()

    cos = torch.cos(angle)
    sin = torch.sin(angle)

    rx, ry, rz = axis.unbind(1)
    zeros = torch.zeros_like(rx)

    K = torch.stack([
        torch.stack([zeros, -rz, ry], dim=1),
        torch.stack([rz, zeros, -rx], dim=1),
        torch.stack([-ry, rx, zeros], dim=1)
    ], dim=2)

    I = torch.eye(3, device=rvecs.device, dtype=rvecs.dtype).expand(K.shape[0], 3, 3)
    R = I + sin.unsqueeze(-1) * K + (1 - cos.unsqueeze(-1)) * torch.bmm(K, K)
    return R


def geodesic_loss(rvec_pred, rvec_true, reduction='mean'):
    R_pred = batch_rodrigues(rvec_pred)
    R_true = batch_rodrigues(rvec_true)
    R_rel = torch.bmm(R_pred.transpose(1, 2), R_true)
    trace = torch.diagonal(R_rel, dim1=-2, dim2=-1).sum(-1)
    cos_angle = (trace - 1) / 2
    cos_angle = torch.clamp(cos_angle, -1 + 1e-6, 1 - 1e-6)
    angle = torch.acos(cos_angle)
    if reduction == 'mean':
        return angle.mean()
    elif reduction == 'none':
        return angle
    else:
        raise NotImplementedError

def pearson_correlation(preds, labels):
    vx = preds - preds.mean(dim=0)
    vy = labels - labels.mean(dim=0)
    corr = (vx * vy).sum(dim=0) / (
        torch.sqrt((vx ** 2).sum(dim=0)) * torch.sqrt((vy ** 2).sum(dim=0)) + 1e-8
    )
    return corr.mean().item()  # 平均所有输出维度的相关系数

# ==========================================================
#                        残差模块（ResBlock）
# ==========================================================
class ResBlock(nn.Module):
    def __init__(self, in_channels):
        super(ResBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(in_channels)

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out += identity
        out = self.relu(out)
        return out


# ==========================================================
#                        模型定义 PoseNet
# ==========================================================
class PoseNet(nn.Module):
    def __init__(self):
        super(PoseNet, self).__init__()
        # Backbone: VGG-style CNN
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(1, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            ResBlock(64),
            nn.MaxPool2d(kernel_size=2, stride=2),

            # Block 2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            # Block 3
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            ResBlock(256),
            nn.MaxPool2d(kernel_size=2, stride=2),

            # Block 4
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            # ResBlock(512),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )

        self.regressor = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512 * 8 * 8, 4096),
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(4096, 2048),
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 3)  # 直接输出原始 rvec（弧度）
        )

    def forward(self, x):
        x = self.features(x)
        rvec = self.regressor(x)
        return rvec


# ==========================================================
#                       自定义数据集类
# ==========================================================
class PoseDataset(Dataset):
    def __init__(self, img_dir, csv_file, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        self.data = []
        with open(csv_file, 'r') as f:
            reader = csv.reader(f)
            next(reader)  # skip header
            for row in reader:
                idx = int(row[0])
                rvec = np.array([float(x) for x in row[1:]], dtype=np.float32)
                self.data.append((idx, rvec))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_idx, rvec = self.data[idx]
        img_path = os.path.join(self.img_dir, f"{img_idx}.png")
        image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        if self.transform:
            image = self.transform(image)
        return image, rvec


# ==========================================================
#           训练主函数（带早停 + 学习率调整）
# ==========================================================
def train_model(img_dir='reproImages',
                csv_file='poseDatas.csv',
                batch_size=32,
                num_epochs=200,
                val_split=0.2,
                patience=25,
                min_delta=0.005,
                save_path='best_pose_model.pth'):

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # 数据转换
    transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((128, 128)),
        transforms.Grayscale(num_output_channels=1),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5], std=[0.5]),
    ])

    # 加载数据集
    dataset = PoseDataset(img_dir=img_dir, csv_file=csv_file, transform=transform)
    val_size = int(val_split * len(dataset))
    train_size = len(dataset) - val_size
    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

    # 初始化模型、损失函数、优化器
    model = PoseNet().to(device)
    criterion = geodesic_loss
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
    scheduler = ReduceLROnPlateau(optimizer,mode='min', factor=0.5, patience=5, threshold=1e-4, threshold_mode='abs', min_lr=1e-6)

    # 早停变量
    best_loss = float('inf')
    trigger_times = 0

    # 用于记录指标
    train_losses, val_losses = [], []
    train_maes, val_maes = [], []
    train_rmses, val_rmses = [], []
    train_corrs, val_corrs = [], []

    # 开始训练
    start_time = time.time()
    for epoch in range(num_epochs):
        # ====== 训练阶段 ======
        model.train()
        running_loss = 0.0
        all_preds, all_labels = [], []

        for images, labels in tqdm(train_loader, desc=f"Epoch {epoch + 1}/{num_epochs} - Training"):
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * images.size(0)
            all_preds.append(outputs.detach().cpu())
            all_labels.append(labels.cpu())


        # 打印当前学习率（可选）
        current_lr = optimizer.param_groups[0]['lr']
        print(f"Epoch [{epoch + 1}/{num_epochs}], Current LR: {current_lr:.2e}")

        epoch_loss = running_loss / len(train_loader.dataset)
        train_losses.append(epoch_loss)

        preds = torch.cat(all_preds)
        labels = torch.cat(all_labels)

        # MAE
        mae = torch.mean(torch.abs(preds - labels)).item()
        train_maes.append(mae)

        # RMSE
        rmse = torch.sqrt(torch.mean((preds - labels) ** 2)).item()
        train_rmses.append(rmse)

        # Correlation
        train_corr = pearson_correlation(preds, labels)
        train_corrs.append(train_corr)

        # ====== 验证阶段 ======
        model.eval()
        val_loss = 0.0
        val_preds, val_labels = [], []

        with torch.no_grad():
            for images, labels in val_loader:
                images = images.to(device)
                labels = labels.to(device)

                outputs = model(images)
                loss = criterion(outputs, labels)

                val_loss += loss.item() * images.size(0)
                val_preds.append(outputs.cpu())
                val_labels.append(labels.cpu())

        val_loss = val_loss / len(val_loader.dataset)
        val_losses.append(val_loss)

        val_preds = torch.cat(val_preds)
        val_labels = torch.cat(val_labels)

        # MAE
        val_mae = torch.mean(torch.abs(val_preds - val_labels)).item()
        val_maes.append(val_mae)

        # RMSE
        val_rmse = torch.sqrt(torch.mean((val_preds - val_labels) ** 2)).item()
        val_rmses.append(val_rmse)

        # Correlation
        val_corr = pearson_correlation(val_preds, val_labels)
        val_corrs.append(val_corr)

        # 打印日志
        print(f"Epoch [{epoch+1}/{num_epochs}], "
              f"Train Loss: {epoch_loss:.6f}, Val Loss: {val_loss:.6f}, "
              f"Train MAE: {mae:.6f}, Val MAE: {val_mae:.6f}, "
              f"Train RMSE: {rmse:.6f}, Val RMSE: {val_rmse:.6f}, "
              f"Train Corr: {train_corr:.6f} "
              f"Val Corr: {val_corr:.6f}")

        # 学习率调度
        scheduler.step(val_loss)

        # 早停机制
        if val_loss < best_loss - min_delta:
            best_loss = val_loss
            trigger_times = 0
            torch.save(model.state_dict(), save_path)
            print("✅ Best model saved.")
        else:
            trigger_times += 1
            print(f"⏳ Early stopping counter: {trigger_times}/{patience}")
            if trigger_times >= patience:
                print("Early stopping triggered.")
                break

    end_time = time.time()
    print(f"Training finished in {end_time - start_time:.2f}s")

    return [train_losses, val_losses, train_maes, val_maes, train_rmses, val_rmses, train_corrs, val_corrs]



if __name__ == "__main__":
    metrics = train_model(
        img_dir='reproImages',
        csv_file='poseDatas.csv',
        batch_size=32,
        num_epochs=200,
        patience=30,
        save_path='best_pose_model.pth'
    )

    train_losses, val_losses, train_maes, val_maes, train_rmses, val_rmses, train_corrs, val_corrs = metrics
    # print(train_losses, val_losses, train_maes, val_maes, train_rmses, val_rmses, train_corrs, val_corrs)
    import matplotlib.pyplot as plt

    plt.figure(figsize=(18, 5))

    plt.subplot(1, 4, 1)
    plt.plot(train_losses, label='Train Loss')
    plt.plot(val_losses, label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss (MSE)')
    plt.legend()

    plt.subplot(1, 4, 2)
    plt.plot(train_maes, label='Train MAE')
    plt.plot(val_maes, label='Validation MAE')
    plt.xlabel('Epoch')
    plt.ylabel('MAE')
    plt.legend()

    plt.subplot(1, 4, 3)
    plt.plot(train_rmses, label='Train RMSE')
    plt.plot(val_rmses, label='Validation RMSE')
    plt.xlabel('Epoch')
    plt.ylabel('RMSE')
    plt.legend()

    plt.subplot(1, 4, 4)
    plt.plot(train_corrs, label='Train Correlation')
    plt.plot(val_corrs, label='Validation Correlation')
    plt.xlabel('Epoch')
    plt.ylabel('Pearson Correlation')
    plt.legend()

    plt.tight_layout()
    plt.show()